// 2D Image Segmentation Example
//
// This example demonstrates how to use GridCut for binary image segmentation.
// We create a simple synthetic image and segment it into foreground and background
// using the graph cut algorithm.

use grid_cut_rs::{GridGraph2D4C, Segment};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("GridCut 2D Image Segmentation Example");
    println!("======================================\n");

    // Image dimensions
    let width = 10;
    let height = 10;
    let num_pixels = width * height;

    println!("Creating {}x{} image...", width, height);

    // Create a synthetic image with a bright square in the center
    let mut image = vec![0.0f32; num_pixels];
    for y in 0..height {
        for x in 0..width {
            let idx = y * width + x;
            // Create a bright square in the center (3x3 region from (3,3) to (6,6))
            if x >= 3 && x < 7 && y >= 3 && y < 7 {
                image[idx] = 1.0; // Bright (foreground)
            } else {
                image[idx] = 0.1; // Dark (background)
            }
        }
    }

    // Print the input image
    println!("\nInput image (0=dark, 1=bright):");
    print_image(&image, width, height);

    // Create the grid graph
    println!("\nSetting up graph cut...");
    let mut graph = GridGraph2D4C::<f32, f32, f32>::new(width, height)?;

    // Define foreground and background models (simple intensity thresholds)
    let foreground_intensity = 1.0;
    let background_intensity = 0.1;

    // Set up terminal capacities (data term)
    // Higher capacity to source = more likely to be foreground
    // Higher capacity to sink = more likely to be background
    let mut cap_source = vec![0.0f32; num_pixels];
    let mut cap_sink = vec![0.0f32; num_pixels];

    for i in 0..num_pixels {
        let intensity = image[i];
        
        // Data term: how well does this pixel match foreground/background?
        let fg_cost = (intensity - foreground_intensity).abs();
        let bg_cost = (intensity - background_intensity).abs();
        
        // Convert costs to capacities (lower cost = higher capacity)
        cap_source[i] = 1.0 / (fg_cost + 0.01); // Add small epsilon to avoid division by zero
        cap_sink[i] = 1.0 / (bg_cost + 0.01);
    }

    // Set up neighbor capacities (smoothness term)
    // Encourage neighboring pixels to have the same label
    let smoothness_weight = 2.0;
    let mut cap_le = vec![0.0f32; num_pixels]; // Left edge
    let mut cap_ge = vec![0.0f32; num_pixels]; // Right edge
    let mut cap_el = vec![0.0f32; num_pixels]; // Up edge
    let mut cap_eg = vec![0.0f32; num_pixels]; // Down edge

    for y in 0..height {
        for x in 0..width {
            let idx = y * width + x;
            let intensity = image[idx];

            // Right neighbor
            if x < width - 1 {
                let neighbor_intensity = image[y * width + (x + 1)];
                let diff = (intensity - neighbor_intensity).abs();
                cap_ge[idx] = smoothness_weight / (diff + 0.1);
            }

            // Down neighbor
            if y < height - 1 {
                let neighbor_intensity = image[(y + 1) * width + x];
                let diff = (intensity - neighbor_intensity).abs();
                cap_eg[idx] = smoothness_weight / (diff + 0.1);
            }

            // Left neighbor (symmetric with right)
            if x > 0 {
                let neighbor_intensity = image[y * width + (x - 1)];
                let diff = (intensity - neighbor_intensity).abs();
                cap_le[idx] = smoothness_weight / (diff + 0.1);
            }

            // Up neighbor (symmetric with down)
            if y > 0 {
                let neighbor_intensity = image[(y - 1) * width + x];
                let diff = (intensity - neighbor_intensity).abs();
                cap_el[idx] = smoothness_weight / (diff + 0.1);
            }
        }
    }

    // Set all capacities at once
    graph.set_caps(&cap_source, &cap_sink, &cap_le, &cap_ge, &cap_el, &cap_eg)?;

    // Compute maximum flow / minimum cut
    println!("Computing maximum flow...");
    graph.compute_maxflow();

    let flow = graph.get_flow();
    println!("Maximum flow: {:.2}", flow);

    // Extract segmentation result
    let mut segmentation = vec![0u8; num_pixels];
    for y in 0..height {
        for x in 0..width {
            let idx = y * width + x;
            let node = graph.node_id(x, y)?;
            let segment = graph.get_segment(node);
            
            segmentation[idx] = match segment {
                Segment::Source => 1, // Foreground
                Segment::Sink => 0,   // Background
            };
        }
    }

    // Print the segmentation result
    println!("\nSegmentation result (0=background, 1=foreground):");
    print_segmentation(&segmentation, width, height);

    // Count foreground and background pixels
    let fg_count = segmentation.iter().filter(|&&x| x == 1).count();
    let bg_count = segmentation.iter().filter(|&&x| x == 0).count();
    println!("\nForeground pixels: {}", fg_count);
    println!("Background pixels: {}", bg_count);

    println!("\n✓ Segmentation complete!");

    Ok(())
}

/// Print a grayscale image as ASCII art
fn print_image(image: &[f32], width: usize, height: usize) {
    for y in 0..height {
        print!("  ");
        for x in 0..width {
            let idx = y * width + x;
            let intensity = image[idx];
            let char = if intensity > 0.5 { '█' } else { '░' };
            print!("{}", char);
        }
        println!();
    }
}

/// Print a binary segmentation as ASCII art
fn print_segmentation(seg: &[u8], width: usize, height: usize) {
    for y in 0..height {
        print!("  ");
        for x in 0..width {
            let idx = y * width + x;
            let char = if seg[idx] == 1 { '█' } else { '░' };
            print!("{}", char);
        }
        println!();
    }
}
